#include "unp.h"
#include <sys/epoll.h>

#define FDSIZE 1024
#define EPOLLEVENTS 20

static void handle_connection(int sockfd);
static void handle_events(int epollfd, struct epoll_event *events, int num, int sockfd, char *buf);
static void do_read(int epollfd, int fd, int sockfd, char *buf);
static void do_write(int epollfd, int fd, int sockfd, char *buf);
static void add_event(int epollfd, int fd, int state);
static void delete_event(int epollfd, int fd, int state);
static void modify_event(int epollfd, int fd, int state);

int main(int argc, char **argv) {
	if (argc != 2) {
		printf("<IPAddress>\n");
		exit(0);
	}
	int sockfd;
	struct sockaddr_in servaddr;
	bzero(&servaddr, sizeof(servaddr));
	sockfd = Socket(AF_INET, SOCK_STREAM, 0);
	servaddr.sin_family = AF_INET;
	servaddr.sin_port = htons(9999);
	inet_pton(AF_INET, argv[1], &servaddr.sin_addr);
	connect(sockfd, (SA*)&servaddr, sizeof(servaddr));
	handle_connection(sockfd);
	close(sockfd);
	return 0;
}

static void handle_connection(int sockfd) {
	int epollfd;
	struct epoll_event events[EPOLLEVENTS];
	char buf[MAXLINE];
	int ret;
	epollfd = epoll_create(FDSIZE);
	add_event(epollfd, STDIN_FILENO, EPOLLIN);
	for (;;) {
		ret = epoll_wait(epollfd, events, EPOLLEVENTS, -1);
		handle_events(epollfd, events, ret, sockfd, buf);
	}
	close(epollfd);
}

static void handle_events(int epollfd, struct epoll_event *events, int num, int sockfd, char *buf) {
	int fd;
	int i;
	for (i = 0; i < num; i++) {
		fd = events[i].data.fd;
		if (events[i].events & EPOLLIN)
			do_read(epollfd, fd, sockfd, buf);
		else if (events[i].events & EPOLLOUT)
			do_write(epollfd, fd, sockfd, buf);
	}
}

static void do_read(int epollfd, int fd, int sockfd, char *buf) {
	int nread;
	nread = read(fd, buf, MAXLINE);
	if (nread == -1) {
		perror("read error");
		close(fd);
	} else if (nread == 0) {
		fprintf(stderr, "server close.\n");
		close(fd);
	} else {
		if (fd == STDIN_FILENO)
			add_event(epollfd, sockfd, EPOLLOUT);
		else {
			delete_event(epollfd, sockfd, EPOLLIN);
			add_event(epollfd, STDOUT_FILENO, EPOLLOUT);
		}
	}
}

static void do_write(int epollfd, int fd, int sockfd, char *buf) {
	int nwrite;
	nwrite = write(fd, buf, strlen(buf));
	if (nwrite == -1) {
		perror("write error:");
		close(fd);
	} else {
		if (fd == STDOUT_FILENO)
			delete_event(epollfd, fd, EPOLLOUT);
		else
			modify_event(epollfd, fd, EPOLLIN);
	}
	memset(buf, 0, MAXLINE);
}

static void add_event(int epollfd, int fd, int state) {
	struct epoll_event ev;
	ev.events = state;
	ev.data.fd = fd;
	epoll_ctl(epollfd, EPOLL_CTL_ADD, fd, &ev);
}

static void delete_event(int epollfd, int fd, int state) {
	struct epoll_event ev;
	ev.events = state;
	ev.data.fd = fd;
	epoll_ctl(epollfd, EPOLL_CTL_DEL, fd, &ev);
}

static void modify_event(int epollfd, int fd, int state) {
	struct epoll_event ev;
	ev.events = state;
	ev.data.fd = fd;
	epoll_ctl(epollfd, EPOLL_CTL_MOD, fd, &ev);
}
	
